package baseFx.database.mybatis;

import baseFx.common.AutoCloseable;
import baseFx.database.Utils;
import org.apache.ibatis.executor.resultset.ResultSetHandler;
import org.apache.ibatis.plugin.*;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.*;

@Intercepts({@Signature(type = ResultSetHandler.class, method = "handleResultSets", args = Statement.class)})
public class MultipleResultInterceptor implements Interceptor {
    private static ThreadLocal<Object> tag = new ThreadLocal<>();

    public AutoCloseable resultSetSpace() {
        tag.set(new Object());
        return () -> tag.remove();
    }

    private ResultSet getFirstResultSet(Statement stmt) throws SQLException {
        ResultSet rs = stmt.getResultSet();
        while (rs == null) {
            if (stmt.getMoreResults()) {
                rs = stmt.getResultSet();
            } else {
                if (stmt.getUpdateCount() == -1) {
                    break;
                }
            }
        }
        return rs;
    }

    private ResultSet getNextResultSet(Statement stmt) {
        try {
            if (stmt.getConnection().getMetaData().supportsMultipleResultSets()) {
                if (!(!stmt.getMoreResults() && stmt.getUpdateCount() == -1)) {
                    ResultSet rs = stmt.getResultSet();
                    if (rs == null) {
                        return getNextResultSet(stmt);
                    } else {
                        return rs;
                    }
                }
            }
        } catch (Exception e) {
            // Intentionally ignored.
        }
        return null;
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        if (tag.get() == null) {
            return invocation.proceed();
        }
        Statement statement = (Statement) invocation.getArgs()[0];
        List<List<Map<?,?>>> data = new ArrayList<>();
        ResultSet rs = getFirstResultSet(statement);
        List<Map<?,?>> rows = Utils.resultSetToList(rs);
        data.add(rows);

        while ((rs = getNextResultSet(statement)) != null) {
            rows = Utils.resultSetToList(rs);
            data.add(rows);
        }
        return data;
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {

    }
}
